# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import torch
from torch.nn import functional as F

from maskrcnn_benchmark.layers import smooth_l1_loss
from maskrcnn_benchmark.modeling.box_coder import BoxCoder
from maskrcnn_benchmark.modeling.matcher import Matcher
from maskrcnn_benchmark.structures.boxlist_ops import boxlist_iou
from maskrcnn_benchmark.modeling.balanced_positive_negative_sampler import (
    BalancedPositiveNegativeSampler
)
from maskrcnn_benchmark.modeling.utils import cat
from maskrcnn_benchmark.structures.bounding_box import BoxList

class FastRCNNLossComputation(object):
    """
    Computes the loss for Faster R-CNN.
    Also supports FPN
    """

    def __init__(
        self,
        proposal_matcher,
        fg_bg_sampler,
        box_coder,
        cls_agnostic_bbox_reg=False,
        allow_empty_gt=True,   # 允许空 GT
    ):
        """
        Arguments:
            proposal_matcher (Matcher)
            fg_bg_sampler (BalancedPositiveNegativeSampler)
            box_coder (BoxCoder)
        """
        self.proposal_matcher = proposal_matcher
        self.fg_bg_sampler = fg_bg_sampler
        self.box_coder = box_coder
        self.cls_agnostic_bbox_reg = cls_agnostic_bbox_reg
        self.allow_empty_gt = allow_empty_gt
        self._sampled_proposals = None
        self._labels = None
        self._regression_targets = None

    def match_targets_to_proposals(self, proposal, target):
        device = proposal.bbox.device
        # ---- 空 GT：直接构造一个与 proposals 对齐的 BoxList，全设为背景 ----
        if len(target) == 0:
            matched_idxs = torch.full((len(proposal),), Matcher.BELOW_LOW_THRESHOLD,  # 纯负样本
                dtype=torch.long, device=device,)
            dummy = BoxList(proposal.bbox.new_zeros((1, 4)), image_size=proposal.size, mode="xyxy")
            matched_targets = dummy[matched_idxs.clamp(min=0)]
            matched_targets.add_field("matched_idxs", matched_idxs)
            return matched_targets
        
        # ---- 正常分支（有 GT）----
        match_quality_matrix = boxlist_iou(target, proposal)
        matched_idxs = self.proposal_matcher(match_quality_matrix)
        # target = target.copy_with_fields("labels")
        matched_targets = target[matched_idxs.clamp(min=0)]
        matched_targets.add_field("matched_idxs", matched_idxs)
        return matched_targets

    def prepare_targets(self, proposals, targets):
        labels = []
        regression_targets = []
        for proposals_per_image, targets_per_image in zip(proposals, targets):
            device = proposals_per_image.bbox.device
            matched_targets = self.match_targets_to_proposals(proposals_per_image, targets_per_image)
            matched_idxs = matched_targets.get_field("matched_idxs")

            if len(targets_per_image) == 0:
                labels_per_image = torch.zeros((len(proposals_per_image),), dtype=torch.int64, device=device)
                regression_targets_per_image = proposals_per_image.bbox.new_zeros((len(proposals_per_image), 4))
            else:
                # 从目标里取出类标
                tgt_labels = targets_per_image.get_field("labels")
                labels_per_image = tgt_labels[matched_idxs.clamp(min=0)].to(torch.int64)
                # IoU 低于下界 => 背景；介于阈值之间 => 忽略
                labels_per_image[matched_idxs == Matcher.BELOW_LOW_THRESHOLD] = 0
                labels_per_image[matched_idxs == Matcher.BETWEEN_THRESHOLDS] = -1
                # 回归目标
                regression_targets_per_image = self.box_coder.encode(matched_targets.bbox, proposals_per_image.bbox)

            labels.append(labels_per_image)
            regression_targets.append(regression_targets_per_image)

        return labels, regression_targets

    def subsample(self, proposals, targets):
        labels, regression_targets = self.prepare_targets(proposals, targets)
        sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels)
        
        sampled = []
        sampled_labels = []
        sampled_regression_targets = []

        for proposals_per_image, labels_per_image, regression_targets_per_image, pos_inds_img, neg_inds_img in zip(
                proposals, labels, regression_targets, sampled_pos_inds, sampled_neg_inds):
            props = proposals_per_image.copy_with_fields([])
            props.add_field("labels", labels_per_image)
            props.add_field("regression_targets", regression_targets_per_image)

            keep = pos_inds_img | neg_inds_img
            props = props[keep]
            sampled.append(props)
            
            sampled_labels.append(labels_per_image[keep])
            sampled_regression_targets.append(regression_targets_per_image[keep])

        self._sampled_proposals = sampled
        self._labels = sampled_labels
        self._regression_targets = sampled_regression_targets
        return sampled

    def __call__(self, class_logits, box_regression):
        assert self._labels is not None and self._regression_targets is not None, \
            "Call loss_evaluator.subsample(...) before loss_evaluator(...)"
        labels = torch.cat(self._labels, dim=0)                    # [M]
        regression_targets = torch.cat(self._regression_targets, dim=0)  # [M, 4]

        # 分类：忽略 -1
        # classification_loss = F.cross_entropy(class_logits, labels, ignore_index=-1, reduction="mean")
        valid = labels != -1
        num_valid = valid.long().sum()

        if num_valid.item() == 0:
            classification_loss = class_logits.sum() * 0.0
        else:
            logits_valid = class_logits[valid]
            labels_valid = labels[valid]

            finite_mask = torch.isfinite(logits_valid)
            if not finite_mask.all():
                logits_valid = torch.where(finite_mask, logits_valid, torch.zeros_like(logits_valid))
            logits_valid = torch.clamp(logits_valid, min=-1e4, max=1e4)
            classification_loss = F.cross_entropy(logits_valid, labels_valid, reduction="mean")

        # ---- 回归损失：只对正样本计算（labels > 0） ----
        sampled_pos_inds_subset = torch.nonzero(labels > 0).squeeze(1)   # [M_pos]
        num_pos = sampled_pos_inds_subset.numel()
        if num_pos == 0:
            box_loss = class_logits.sum() * 0.0
            return classification_loss, box_loss
        
        # 取出正样本的目标回归值
        regression_targets_pos = regression_targets[sampled_pos_inds_subset]  # [M_pos, 4]
        if self.cls_agnostic_bbox_reg:
            box_regression_pos = box_regression[sampled_pos_inds_subset, :]   # [M_pos, 4]
        else:
            # 每个正样本取其 GT 类别对应的 4-d 通道
            labels_pos = labels[sampled_pos_inds_subset]                       # [M_pos]
            box_regression_pos = box_regression.view(box_regression.size(0), -1, 4)  # [M, K, 4]
            box_regression_pos = box_regression_pos[sampled_pos_inds_subset, labels_pos]  # [M_pos, 4]

        # SmoothL1：用 sum 再除以正样本数，返回标量
        box_loss = smooth_l1_loss(
            box_regression_pos,
            regression_targets_pos,
            beta=1.0 / 9,
            reduction="sum"
        ) / float(num_pos)

        return classification_loss, box_loss


def make_roi_box_loss_evaluator(cfg):
    matcher = Matcher(
        cfg.MODEL.ROI_HEADS.FG_IOU_THRESHOLD,
        cfg.MODEL.ROI_HEADS.BG_IOU_THRESHOLD,
        allow_low_quality_matches=False,
    )

    bbox_reg_weights = cfg.MODEL.ROI_HEADS.BBOX_REG_WEIGHTS
    box_coder = BoxCoder(weights=bbox_reg_weights)

    fg_bg_sampler = BalancedPositiveNegativeSampler(
        cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE, cfg.MODEL.ROI_HEADS.POSITIVE_FRACTION
    )

    cls_agnostic_bbox_reg = cfg.MODEL.CLS_AGNOSTIC_BBOX_REG

    loss_evaluator = FastRCNNLossComputation(
        matcher,
        fg_bg_sampler,
        box_coder,
        cls_agnostic_bbox_reg
    )

    return loss_evaluator
